#include "codeml.h"
#include "debug_func.h"
#include <errno.h>
#include <stdexcept>
#include <algorithm>
#include <iostream>
#include <sstream>
#include <iostream>
#include <fstream>
#include <vector>
#include <boost/lexical_cast.hpp>
#include <sys/stat.h>
#include <unistd.h>
#include <sys/types.h>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/regex.hpp>
#include <ultoa.h>
#include <pstream.h>
using std::vector;
using std::replace;
using std::ofstream;
using std::ifstream;
using std::string;
using std::cerr;
using std::istringstream;
using boost::lexical_cast;
using boost::filesystem::path;





namespace
{
void parse_codeml_results(	std::vector<t_codeml_results>&		results,
							unsigned							seq_tag1,
							unsigned							seq_tag2,
							const std::string&					error_string,
							const std::string&					file_name = "results");
void parse_yn00_results (	std::vector<t_codeml_results>&		results,
							unsigned							seq_tag1,
							unsigned							seq_tag2,
							const std::string&					error_string,
							const std::string&					file_name = "results");



void parse_yn00_results(vector<t_codeml_results>&		results,
						unsigned						seq_tag1,
						unsigned						seq_tag2,
						const string&					error_string,
						const std::string&				file_name)
{
	try
	{
		// get results
		ifstream f_results (file_name.c_str());
		if (!f_results)
			throw std::runtime_error("Could not open PAML results file [" + file_name +"] "
								" probably because the programme did not run successfully.");
	
		// get last line
		string new_line, line;
		unsigned line_num = 0;
		while (getline(f_results, new_line))
		{
			++line_num;
			line = new_line;
		}
	
		replace(line.begin(), line.end(), '+', ' ');
		replace(line.begin(), line.end(), '-', ' ');
		istringstream istrm(line);
		string str_discard;
		float s, n, ka, ks, ka_se, ks_se, kaks, kappa, dT;
	//	seq. seq.     S       N        t   kappa   omega     dN +- SE    dS +- SE
	//   2    1   107.9   486.1   6.8501  2.0675  0.5728 2.0110 +- 0.4455  3.5107 +- 4.5383
	
		istrm >> str_discard >> str_discard >> s >> n >>
				dT >> kappa >> kaks >>
				ka >> ka_se  >> ks >> ks_se ;
		if (istrm.fail())
		{
			throw std::runtime_error("The Kaks value could not be parsed.");
			//cerr << "Failed at line " << line_num << "\n[" << line << "]\n";
		}
		results.push_back(	t_codeml_results(	seq_tag1, seq_tag2, n, s, ka, ks, kaks, 0.0,
												"", ka_se, ks_se, kappa, dT));

	}
	catch(const std::runtime_error& e)
	{
		results.push_back(	t_codeml_results(seq_tag1, seq_tag2, "Error:" + 
												error_string + " " + e.what()));
	}
}

void parse_codeml_results(vector<t_codeml_results>&		results,
						unsigned						seq_tag1,
						unsigned						seq_tag2,
						const string&					error_string,
						const std::string&				file_name)
{

	try
	{
		// get results
		ifstream f_results (file_name.c_str());
		if (!f_results)
		{
			throw std::runtime_error("Could not open PAML results file [" + file_name +"] "
								" probably because the programme did not run successfully.");
		}

		string new_line, line;
		//
		// get LnL
		//
		float lnL = 99.0;
		string saved_line;
		unsigned cnt_lines = 0;
		while (getline(f_results, new_line))
		{
			if (!new_line.length())
				continue;
			if(boost::find_first(new_line, "seed used = "))
				continue;

			++cnt_lines;
			// find lnL followed by =
			if (new_line.compare(0,  3,  "lnL") != 0)
				continue;
			saved_line = new_line;
			string::size_type pos = new_line.find('=', 3);
			if (pos == string::npos)
				continue;
			istringstream is (new_line.substr(pos+1));
			is >> lnL;
			if (is.fail())
				throw std::runtime_error("Could not parse lnL from [" + new_line + "]");
			break;
		}

		if (!f_results)
		{
			if (cnt_lines)
				throw std::runtime_error("The line lnL= could not be found in [" + saved_line + "]");
			else
				throw std::runtime_error("Could not open PAML results file [" + file_name +"] "
									" probably because the programme did not run successfully.");

		}


		//
		// get kappa
		//
		float kappa;
		float float_discard;
		f_results >> float_discard >> kappa;
		if (!f_results)
			throw std::runtime_error("The value of kappa could not be parsed.");

		while (getline(f_results, new_line))
		{
			if (new_line.compare(0,  2,  "t=") == 0)
				break;
		}
		if (!f_results)
			throw std::runtime_error("The line containing t, omega etc. could not be found.");

		//
		// get KaKs
		//
		float s, n, ka, ks, ka_se, ks_se, kaks, dT;
		{
			//	t= 0.0313  S=   111.8  N=   368.2  dN/dS= 2.3030  dN= 0.0120  dS= 0.0052
			replace(new_line.begin(), new_line.end(), '=', ' ');
			istringstream istrm(new_line);
			string str_discard;
			istrm >> 	str_discard >> dT				>>		// t= 0.0313
						str_discard >> s				>>		// S=   111.8
						str_discard >> n				>>		// N=   368.2
						str_discard >> kaks				>>		// dN/dS= 2.3030
						str_discard >> ka				>>		// dN= 0.0120
						str_discard >> ks;						// dS= 0.0052
			if (istrm.fail())
				throw std::runtime_error("The StdErr line could not be parsed [" + new_line + "].");
		}

		//
		// Stderr
		//
        if (!getline(f_results, new_line))
			throw std::runtime_error("The StdErr line could not be parsed.");
		{
			replace(new_line.begin(), new_line.end(), '=', ' ');
			replace(new_line.begin(), new_line.end(), '+', ' ');
			replace(new_line.begin(), new_line.end(), '-', ' ');
			istringstream istrm(new_line);
			string str_discard;
			float float_discard;
			istrm	>> str_discard >> float_discard >> ka_se
					>> str_discard >> float_discard >> ks_se;
			if (istrm.fail())
				throw std::runtime_error("The StdErr line could not be parsed [" + new_line + "].");
		}

		results.push_back(	t_codeml_results(	seq_tag1, seq_tag2, n, s, ka, ks, kaks, lnL,
												"", ka_se, ks_se, kappa, dT));
	}
	catch(const std::runtime_error& e)
	{
		results.push_back(	t_codeml_results(seq_tag1, seq_tag2, string("Error:") + 
												e.what() + " [" + error_string + ']'));
	}
}






unsigned remove_stop_codons(string& seq)
{
	unsigned cnt_stops = 0;
	string::iterator i = seq.begin();
	string::iterator end = seq.end();
	for (;i != end; i += 3)
	{
		if (*i != 'T')
			continue;
		if	(
				(*(i+1) == 'A' && 	(*(i+2) == 'A' 	||
									 *(i+2) == 'G')	)
				||
				(*(i+1) == 'G' &&    *(i+2) == 'A')
			)
		{
			*(i  ) = '-';
			*(i+1) = '-';
			*(i+2) = '-';
			++cnt_stops;
		}
	}
	return cnt_stops;
}


//________________________________________________________________________________________

//	kaks_analyse

//________________________________________________________________________________________
void retrieve_paml_stdout_stderr(const string& exe_path, redi::pstream& rs, string& error_string)
{
	// get stdin
	unsigned cnt_lines = 0;
	{
		++cnt_lines;
		string line;
		while (getline(rs.err(), line))
			error_string += line + "\n";
		rs.clear();
	}

	// if command not found, throw here
	if(cnt_lines == 1 && error_string.substr(0, 4) == "sh: ")
		 throw std::runtime_error("Could not find [" + exe_path + "]\n" + 
					error_string.substr(4));

	// get stdout
	// get stdin
	{
		string line;
		while (getline(rs.out(), line))
			error_string += line + "\n";
		rs.clear();
	}
}

}
















namespace BioCPP
{
bool kaks_analyse (	string 							seq1,
					string		 					seq2,
					unsigned						seq_tag1,
					unsigned						seq_tag2,
					vector<t_codeml_results>&   	results,
					const string&					kaks_exe_path,
					BioCPP::eDNDS_PROGRAM			dnds_program,
					unsigned&						cnt_hung_paml_restarted)
{
	if (!seq1.length() % 3 != 0)
		throw std::logic_error("The sequence \n[" + seq1 + "]\n contained a non-integral "
								"number of codons.");
	if (!seq2.length() % 3 != 0)
		throw std::logic_error("The sequence \n[" + seq2 + "]\n contained a non-integral "
								"number of codons.");

	string extra_error_string;
	if (unsigned cnt_stops = remove_stop_codons(seq1))
		extra_error_string += ultoa(cnt_stops) + " stop codons removed from sequence1. ";
	if (unsigned cnt_stops = remove_stop_codons(seq2))                                 
		extra_error_string += ultoa(cnt_stops) + " stop codons removed from sequence2. ";

	// remove any existing results file before running analysis
	remove( "results");

	//
	//	sequences file
	//
	ofstream seq_strm ("seq.tmp");
	if (!seq_strm)
		throw std::runtime_error("Could not open seq.tmp, the sequences file for codeml.");;
	seq_strm << "2 " << seq1.length() << "\n";
	seq_strm << "a  " << seq1 << "\n";
	seq_strm << "b  " << seq2 << "\n";
	seq_strm.close();

	string error_string;
	{
		bool success = false;

		// retry if hung
		unsigned tries = 0;
		for (; tries < 3; tries++)
		{
			redi::pstream rs(kaks_exe_path, redi::pstreams::pstdout | redi::pstreams::pstderr);

			// poll for 5 minutes
			for (unsigned i = 0; i < 1200; ++i)
			{
				// poll every 0.25 s
				usleep(250000);
				if (rs.rdbuf()->peek_exited())
				{
					success = true;
					retrieve_paml_stdout_stderr(kaks_exe_path, rs, error_string);
					break;
				}
			}
			if (success)
				break;
			else
				++cnt_hung_paml_restarted;

			// kill child and restart
			if(!rs.rdbuf()->kill(9))
			{
				char buf [ 2048];
				strerror_r(rs.rdbuf()->error(), buf, 2047);
				throw std::runtime_error(string("Failed to kill hung PAML child process\n")
										  + buf);
			}
		}
		
		if (!success)
			return false;


		// log errors
		replace(error_string.begin(), error_string.end(), '\t', ' ');
		replace(error_string.begin(), error_string.end(), '\\', ' ');
		if(boost::find_first(error_string, "var(dS,dN) not calculable"))
			extra_error_string += "The covariance matrix for dS, dN could not be "
										"calculated. Some parameters may be out of range.";
		//cerr << error_string;
	}
	if (dnds_program == BioCPP::eCODEML)
		parse_codeml_results(results, seq_tag1, seq_tag2, error_string);
	else
		parse_yn00_results(results, seq_tag1, seq_tag2, error_string);

	unlink("results");
	
	// add error s
	if (extra_error_string.length())
		results.back().error_str += extra_error_string;

	return true;
}

//________________________________________________________________________________________

//	prepare_temp_directory

//		do kaks in temp directory to avoid crud

//________________________________________________________________________________________
void prepare_temp_directory(const string&			tmp_file_path,
							bool					fixed_kappa,
							float					kappa,
							BioCPP::eDNDS_PROGRAM	dnds_program,
							BioCPP::eCODEML_CODON	f61,
							bool&					directory_created)
{
	debug_func_cerr(5);

	// create temp directory if necessary
	directory_created = false;
	if (!boost::filesystem::exists(path(tmp_file_path)))
	{
		if (::mkdir( tmp_file_path.c_str(), S_IRWXU|S_IRWXG|S_IRWXO ) != 0 )
			throw std::runtime_error("Could not create " + tmp_file_path);

		// successively created
		directory_created = true;
	}
	else
	{
		if(!boost::filesystem::is_directory(path(tmp_file_path)))
			throw std::runtime_error(tmp_file_path + " already exists as a file");
	}

	// change to temp directory
	if (chdir(tmp_file_path.c_str()))
	{
		char buf[256];
		buf[255] = '\0';
		char * error_str = strerror_r(errno, buf, 255);
		throw std::runtime_error("Could not change to " + tmp_file_path + "[ errno = "  +
									lexical_cast<string>(errno) + " ]\n" +
									(error_str ? error_str : ""));
	}

	// write contro file for kaks
	ofstream f_kaksctl("kaks.ctl");
	if (!f_kaksctl)
		throw std::runtime_error("Could not open " + tmp_file_path + "kaks.ctl");

	f_kaksctl <<  	"seqfile = seq.tmp\n"
					"outfile = results\n"
					"icode = 0\n";
	if (dnds_program == BioCPP::eCODEML)
	{
		f_kaksctl <<"noisy = 0\n"
					"verbose = 0\n"
					"runmode = -2\n"
					"seqtype = 1\n"
					"CodonFreq = "
				<< ( f61 == eF61 ? 3 : 2)
				<<
					"\n"  
					"aaDist = 0\n"
					"model = 0\n"
					"NSsites = 0\n"
					"Mgene = 0\n"
					"fix_omega = 0\n"
					"omega = .4\n"
					"fix_alpha = 1\n"
					"alpha = 0.\n"
					"Malpha = 0\n"
					"ncatG = 8\n"
					"clock = 0\n"
					"getSE = 1\n"
					"RateAncestor = 1\n"
					"Small_Diff = .5e-7\n"
					"cleandata = 1\n"
					"method = 0\n";
		if (fixed_kappa)
			f_kaksctl <<
					"fix_kappa = 1\n"
					"kappa = " << kappa << "\n";
		else
			f_kaksctl <<
					"fix_kappa = 0\n"
					"kappa = 2\n";

	}
	else
		f_kaksctl <<
//					"verbose = 0\n"
//					"noisy = 0\n"
					"weighting = 1\n"
					"commonf3x4 = 0\n";

	f_kaksctl.close();
}

};
